import torchvision.transforms as T
import random
import numpy as np
from PIL import Image
import torch


class Compose(object):
    """组合多种数据增强操作"""
    
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img, mask):
        for t in self.transforms:
            img, mask = t(img, mask)
        return img, mask


class RandomHorizontalFlip(object):
    """随机水平翻转"""
    
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img, mask):
        if random.random() < self.p:
            img = T.functional.hflip(img)
            mask = T.functional.hflip(mask)
        return img, mask


class RandomRotation(object):
    """随机旋转"""
    
    def __init__(self, degrees):
        self.degrees = degrees

    def __call__(self, img, mask):
        angle = T.RandomRotation.get_params((-self.degrees, self.degrees))
        img = T.functional.rotate(img, angle, interpolation=T.InterpolationMode.BILINEAR)
        mask = T.functional.rotate(mask, angle, interpolation=T.InterpolationMode.NEAREST)
        return img, mask


class RandomRescale(object):
    """随机缩放"""
    
    def __init__(self, scale_range):
        self.scale_range = scale_range

    def __call__(self, img, mask):
        scale_factor = random.uniform(self.scale_range[0], self.scale_range[1])
        new_size = (int(img.size[0] * scale_factor), int(img.size[1] * scale_factor))
        img = T.functional.resize(img, new_size, interpolation=T.InterpolationMode.BILINEAR)
        mask = T.functional.resize(mask, new_size, interpolation=T.InterpolationMode.NEAREST)
        return img, mask


class Resize(object):
    """调整图像和掩码到指定大小"""
    
    def __init__(self, size):
        self.size = size

    def __call__(self, img, mask):
        img = T.functional.resize(img, self.size, interpolation=T.InterpolationMode.BILINEAR)
        mask = T.functional.resize(mask, self.size, interpolation=T.InterpolationMode.NEAREST)
        return img, mask


class ToTensor(object):
    """将PIL图像转换为Tensor"""
    
    def __call__(self, img, mask):
        img = T.functional.to_tensor(img)
        mask = torch.from_numpy(np.array(mask, dtype=np.int64))
        return img, mask


class Normalize(object):
    """归一化图像"""
    
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, img, mask):
        img = T.functional.normalize(img, mean=self.mean, std=self.std)
        return img, mask


def get_train_augmentation(config, current_epoch=0):
    """根据配置和当前epoch获取训练数据增强变换"""
    
    transforms = []
    if config["augmentations"]["enabled"] and current_epoch < config["augmentations"]["augment_stop_epoch"]:
        if config["augmentations"]["flip"]:
            transforms.append(RandomHorizontalFlip())
        if config["augmentations"]["rotate"]:
            transforms.append(RandomRotation(config["augmentations"]["rotate"]))
        if config["augmentations"]["scale"]:
            transforms.append(RandomRescale(config["augmentations"]["scale"]))
    
    transforms.append(Resize(config["data"]["image_size"]))
    transforms.append(ToTensor())
    transforms.append(Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])) # ImageNet均值和标准差
    
    return Compose(transforms)


def get_val_augmentation(config):
    """获取验证数据增强变换"""
    
    transforms = [
        Resize(config["data"]["image_size"]),
        ToTensor(),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]
    return Compose(transforms)


def get_inference_augmentation(config):
    """获取推理数据增强变换"""
    
    transforms = [
        Resize(config["data"]["image_size"]),
        ToTensor(),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]
    return Compose(transforms)


if __name__ == "__main__":
    # 简单测试
    from PIL import Image
    import torch
    import yaml
    
    # 创建一个假的config.yaml内容
    fake_config_content = """
    data:
      image_size: [256, 256]
    augmentations:
      enabled: True
      augment_stop_epoch: 5
      flip: True
      rotate: 10
      scale: [0.8, 1.2]
    """
    
    config = yaml.safe_load(fake_config_content)
    
    # 创建一个假的图像和掩码
    dummy_img = Image.new("RGB", (512, 512), color = 'red')
    dummy_mask = Image.new("L", (512, 512), color = 1) # 假设单通道掩码
    
    print("原始图像大小:", dummy_img.size)
    
    # 测试训练增强
    train_transform = get_train_augmentation(config, current_epoch=0)
    transformed_img, transformed_mask = train_transform(dummy_img, dummy_mask)
    print("训练增强后图像形状:", transformed_img.shape)
    print("训练增强后掩码形状:", transformed_mask.shape)
    
    # 测试验证增强
    val_transform = get_val_augmentation(config)
    transformed_img_val, transformed_mask_val = val_transform(dummy_img, dummy_mask)
    print("验证增强后图像形状:", transformed_img_val.shape)
    print("验证增强后掩码形状:", transformed_mask_val.shape)
    
    # 测试推理增强
    inference_transform = get_inference_augmentation(config)
    transformed_img_inf, transformed_mask_inf = inference_transform(dummy_img, dummy_mask)
    print("推理增强后图像形状:", transformed_img_inf.shape)
    print("推理增强后掩码形状:", transformed_mask_inf.shape)


